#include "util.h"
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <stdio.h>
#include <unistd.h>
#include "des.h"

#define BUFFER_SIZE 4096

octet_string new_octet_string() {
    octet_string rt;
    rt.size = 0;
    rt.data = NULL;
    return rt;
}

multi_msg new_multi_msg() {
    multi_msg rt;
    rt.num = 0;
    rt.msgs = NULL;
    return rt;
}

uint64_t rand_uint64() {
    uint64_t rt = 0;
    for (int i = 0; i < 64; ++i) {
        rt |= (((uint64_t)(rand() & 1)) << i);
    }
    return rt;
}

void mm_dumps(octet_string *os, const multi_msg *m) {
    if (os->size) {
        os_free(os);
    }
    os->size = 1 + m->num * 2;
    for (size_t i = 0; i < m->num; ++i) {
        os->size += m->msgs[i].size;
    }
    os->data = (uint8_t*)malloc(os->size);
    uint8_t *a = os->data;
    a[0] = (uint8_t)m->num;
    int cur = 1;
    for (size_t i = 0; i < m->num; ++i) {
        const single_msg *msg = m->msgs + i;
        a[cur++] = msg->size & 0xFF;
        a[cur++] = (msg->size >> 8) & 0xFF;
        memcpy(a + cur, msg->data, msg->size);
        cur += msg->size;
    }
}

void mm_loads(multi_msg *m, const octet_string *os) {
    if (m->num != 0) {
        mm_free(m);
    }
    const uint8_t *str = os->data;
    m->num = str[0];
    m->msgs = (single_msg*)malloc(sizeof(single_msg) * m->num);
    int cur = 1;
    for (size_t i = 0; i < m->num; ++i) {
        single_msg *msg = m->msgs + i;
        msg->size = str[cur++];
        msg->size |= (((size_t)str[cur++]) << 8);
        msg->data = (uint8_t*)malloc(msg->size);
        memcpy(msg->data, str + cur, msg->size);
        cur += msg->size;
    }
}

void mm_build(multi_msg *m, const char *format, ...) {
    m->num = 0;
    for (int i = 0; format[i]; ++i) {
        if (format[i] == 's' || format[i] == 'd' || format[i] == 'k' || format[i] == 'o') {
            ++m->num;
        }
    }
    m->msgs = (single_msg*)malloc(sizeof(single_msg) * m->num);

    va_list valist;
    va_start(valist, format);
    int cnt = 0;
    for (int i = 0; format[i]; ++i) {
        single_msg *msg = m->msgs + cnt;
        switch (format[i]) {
            case 's': {
                const char *s = va_arg(valist, char*);
                msg->size = strlen(s) + 1;
                msg->data = (uint8_t*)malloc(msg->size);
                memcpy(msg->data, s, msg->size);
                break;
            }
            case 'd': {
                int32_t d = va_arg(valist, int32_t);
                msg->size = 4;
                msg->data = (uint8_t*)malloc(msg->size);
                const uint32_t *p = (uint32_t*)(&d);
                for (int i = 0; i < 4; ++i) {
                    msg->data[i] = (((*p) >> (i * 8)) & 0xFF);
                }
                break;
            }
            case 'k': {
                uint64_t k = va_arg(valist, uint64_t);
                msg->size = 8;
                msg->data = (uint8_t*)malloc(msg->size);
                for (int i = 0; i < 8; ++i) {
                    msg->data[i] = ((k >> (i * 8)) & 0xFF);
                }
                break;
            }
            case 'o': {
                const octet_string *os = va_arg(valist, octet_string*);
                msg->size = os->size;
                msg->data = (uint8_t*)malloc(msg->size);
                memcpy(msg->data, os->data, msg->size);
                break;
            }
            default:
                --cnt;
        }
        ++cnt;
    }
    va_end(valist);
}

void mm_free(multi_msg *m) {
    for (size_t i = 0; i < m->num; ++i) {
        free(m->msgs[i].data);
    }
    free(m->msgs);
    m->msgs = NULL;
    m->num = 0;
}

void os_free(octet_string *os) {
    free(os->data);
    os->size = 0;
}

void mm_print(const multi_msg *m) {
    printf("tot %zu messages\n", m->num);
    for (size_t i = 0; i < m->num; ++i) {
        printf("[%zu] (%zu)", i, m->msgs[i].size);
        for (size_t j = 0; j < m->msgs[i].size; ++j) {
            printf(" %02X", m->msgs[i].data[j]);
        }
        putchar('\n');
    }
    putchar('\n');
}

void mm_sock_write(int fd, const multi_msg *m) {
    octet_string os = new_octet_string();
    mm_dumps(&os, m);
    write(fd, os.data, os.size);
    os_free(&os);
}

void mm_sock_read(multi_msg *m, int fd) {
    static uint8_t recvBuffer[BUFFER_SIZE];
    octet_string os;
    os.size = read(fd, recvBuffer, sizeof(recvBuffer));
    os.data = recvBuffer;
    mm_loads(m, &os);
}

void mm_dumps_encrypt(octet_string *os, const multi_msg *m, uint64_t key) {
    octet_string plain = new_octet_string();
    mm_dumps(&plain, m);
    if (os->size) {
        os_free(os);
    }
    os->size = get_number_of_encrypted_bytes(plain.size);
    os->data = (uint8_t*)malloc(os->size);
    des_init(key);
    des_encrypt_bytes(os->data, plain.data, plain.size);
    os_free(&plain);
}

void mm_decrypt_loads(multi_msg *m, const octet_string *os, uint64_t key) {
    octet_string plain = new_octet_string();
    plain.data = (uint8_t*)malloc(os->size);
    des_init(key);
    des_decrypt_bytes(plain.data, &plain.size, os->data, os->size);
    mm_loads(m, &plain);
    os_free(&plain);
}
